# %%
import torch
import torchvision
import torchvision.transforms as transforms
import torchvision.datasets as datasets

import matplotlib.pyplot as plt
from pathlib import Path


DATA_PATH = Path(
    r'C:\files\git_repository\pytorch-learning\datasets\hymenoptera_data\train')
# %%
train_transform = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0), ratio=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    torchvision.transforms.ColorJitter(
        brightness=0.5, contrast=0, saturation=0, hue=0),
    torchvision.transforms.ColorJitter(
        brightness=0, contrast=0.5, saturation=0, hue=0),
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
data = datasets.ImageFolder(DATA_PATH, transform=train_transform)
# %%
data_loader = torch.utils.data.DataLoader(data, batch_size=32, shuffle=True)
#%%
x,y = next(iter(data_loader))
print(x.shape)
print(y.shape)
#%%
class_name_dict = data.class_to_idx
print(class_name_dict)
class_name_dict = dict(zip(class_name_dict.values(),class_name_dict.keys()))
print(class_name_dict)
#%%
import matplotlib.pyplot as plt

for i in range(3):
    y_item = y[i].item()
    name = class_name_dict[y_item]
    plt.subplot(1,3,i+1).set_title(name)
    x_item = x[i]
    x_item = x_item.permute(1,2,0)
    # print(x_item.shape)
    plt.imshow(x_item)
# plt.imshow(x)
# %%

# 指定读取的图片路径
# transform函数组合
train_transform = transforms.Compose([
    transforms.Resize(224),
    transforms.RandomResizedCrop(224, scale=(0.6, 1.0), ratio=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    torchvision.transforms.ColorJitter(
        brightness=0.5, contrast=0, saturation=0, hue=0),
    torchvision.transforms.ColorJitter(
        brightness=0, contrast=0.5, saturation=0, hue=0),
    transforms.ToTensor(),
    transforms.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])


# 使用ImageFolder读取数据
all_data = torchvision.datasets.ImageFolder(
    root=DATA_PATH,
    transform=train_transform
)


# 定义数据加载器
train_set = torch.utils.data.DataLoader(
    all_data,
    batch_size=32,
    shuffle=True
)
